package com.sicau.utils;

import com.itextpdf.kernel.pdf.PdfDocument;
import com.itextpdf.kernel.pdf.PdfPage;
import com.itextpdf.kernel.pdf.PdfReader;
import com.itextpdf.kernel.pdf.canvas.parser.EventType;
import com.itextpdf.kernel.pdf.canvas.parser.PdfCanvasProcessor;
import com.itextpdf.kernel.pdf.canvas.parser.data.IEventData;
import com.itextpdf.kernel.pdf.canvas.parser.listener.IEventListener;
import com.itextpdf.kernel.pdf.canvas.parser.listener.SimpleTextExtractionStrategy;
import com.itextpdf.kernel.pdf.xobject.PdfImageXObject;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.stereotype.Component;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * PDF 文本+图片内容提取工具（iText 9 版本）
 * 功能：提取文本 + 提取图片并通过 ChatClient 识别内容，拼接成完整文本
 */
@Component
public class PdfTextAndImageExtractor {

    private final ChatClient chatClient;

    public PdfTextAndImageExtractor(ChatModel chatModel) {
        chatClient = ChatClient
                .builder(chatModel)
                .defaultSystem("请识别图片内容")
                .build();
    }

    /**
     * 提取 PDF 所有页面的文本+图片识别内容（拼接结果）
     * @param fileInputStream PDF 文件输入流（如 MultipartFile 的 getInputStream()）
     * @return 拼接后的完整文本（文本内容 + 图片识别内容）
     */
    public String extractAllContent(InputStream fileInputStream) {
        PdfReader pdfReader = null;
        PdfDocument pdfDocument = null;
        try {
            // 1. 初始化 PDF 读取器和文档对象
            pdfReader = new PdfReader(fileInputStream);
            pdfDocument = new PdfDocument(pdfReader);

            StringBuilder fullContent = new StringBuilder();

            // 2. 遍历 PDF 所有页面（若只需第一页，可改为 pdfDocument.getFirstPage()）
            for (int pageNum = 1; pageNum <= pdfDocument.getNumberOfPages(); pageNum++) {
                PdfPage page = pdfDocument.getPage(pageNum);
                fullContent.append("=== 第 ").append(pageNum).append(" 页内容 ===\n");

                // 3. 提取当前页面的文本 + 图片识别内容
                String pageContent = extractPageContent(page);
                fullContent.append(pageContent).append("\n\n");
            }

            return fullContent.toString().trim();

        } catch (Exception e) {
            e.printStackTrace();
            return "PDF 解析失败：" + e.getMessage();
        } finally {
            // 4. 关闭资源（避免内存泄漏）
            if (pdfDocument != null) {
                pdfDocument.close();
            }
            if (pdfReader != null) {
                try {
                    pdfReader.close();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }
    }

    /**
     * 提取单个 PDF 页面的文本 + 图片识别内容
     * @param page PDF 页面对象
     * @return 单个页面的完整内容（文本 + 图片识别结果）
     */
    private String extractPageContent(PdfPage page) throws Exception {
        // 存储页面中的图片信息（图片对象 + 位置信息，用于排序）
        List<ImageWithPosition> imageList = new ArrayList<>();

        // 5. 自定义监听器：同时提取文本和图片
        CustomContentListener listener = new CustomContentListener(imageList);
        PdfCanvasProcessor processor = new PdfCanvasProcessor(listener);
        processor.processPageContent(page);

        // 6. 提取页面文本
        String textContent = listener.getTextContent();

        // 7. 处理图片：调用 ChatClient 识别图片内容
        StringBuilder imageOcrContent = new StringBuilder();
        for (int i = 0; i < imageList.size(); i++) {
            ImageWithPosition imageWithPos = imageList.get(i);
            PdfImageXObject imageXObject = imageWithPos.imageXObject;

            // 7.1 将 PdfImageXObject 转为 InputStream（传给 ChatClient）
            byte[] imageBytes = imageXObject.getImageBytes();
            InputStream imageInputStream = new ByteArrayInputStream(imageBytes);

            // 7.2 调用 ChatClient 识别图片内容（核心：替换为你的实际识别逻辑）

            String ocrResult = chatClient.prompt()
                    .messages(new UserMessage(new ByteArrayResource(imageBytes)))
                    .call()
                    .content();

            // 7.3 拼接图片识别结果（标记图片序号和位置）
            imageOcrContent.append("\n【图片").append(i + 1).append("识别结果】：")
                    .append(ocrResult == null ? "图片识别失败" : ocrResult)
                    .append("\n");

            // 关闭图片流
            imageInputStream.close();
        }

        // 8. 拼接文本和图片识别结果（文本在前，图片识别结果在后；可根据位置调整顺序）
        return textContent + imageOcrContent.toString();
    }

    /**
     * 自定义监听器：同时提取文本和图片（实现 IEventListener 接口）
     */
    private static class CustomContentListener implements IEventListener {
        // 文本提取策略（复用 iText 原生策略）
        private final SimpleTextExtractionStrategy textStrategy = new SimpleTextExtractionStrategy();
        // 存储图片及位置信息
        private final List<ImageWithPosition> imageList;

        public CustomContentListener(List<ImageWithPosition> imageList) {
            this.imageList = imageList;
        }

        // 获取提取的文本内容
        public String getTextContent() {
            return textStrategy.getResultantText();
        }

        @Override
        public void eventOccurred(IEventData data, EventType type) {
            // 9. 处理文本事件：委托给 SimpleTextExtractionStrategy
            if (type.equals(EventType.RENDER_TEXT)) {
                textStrategy.eventOccurred(data, type);
            }
            // 10. 处理图片事件：捕获 PdfImageXObject 并存储
            else if (type.equals(EventType.RENDER_IMAGE)) {
                // 解析图片数据
                com.itextpdf.kernel.pdf.canvas.parser.data.ImageRenderInfo imageRenderInfo =
                        (com.itextpdf.kernel.pdf.canvas.parser.data.ImageRenderInfo) data;
                PdfImageXObject imageXObject = imageRenderInfo.getImage();

                if (imageXObject != null) {
                    // 存储图片对象和位置信息（位置信息可用于调整识别结果顺序）
                    float x = imageRenderInfo.getImageCtm().get(6); // 图片左下角 X 坐标
                    float y = imageRenderInfo.getImageCtm().get(7); // 图片左下角 Y 坐标
                    imageList.add(new ImageWithPosition(imageXObject, x, y));
                }
            }
        }

        // 声明支持的事件类型（文本 + 图片）
        @Override
        public Set<EventType> getSupportedEvents() {
            Set<EventType> events = new HashSet<>();
            events.add(EventType.RENDER_TEXT);
            events.add(EventType.RENDER_IMAGE);
            return events;
        }
    }

    /**
     * 辅助类：存储图片对象和位置信息（用于排序或标记）
     */
    private static class ImageWithPosition {
        PdfImageXObject imageXObject;
        float x; // 图片 X 坐标（PDF 坐标系：左下角为原点）
        float y; // 图片 Y 坐标

        public ImageWithPosition(PdfImageXObject imageXObject, float x, float y) {
            this.imageXObject = imageXObject;
            this.x = x;
            this.y = y;
        }
    }

}